import tkinter as tk
import tensorflow as tf
from PIL import Image, ImageGrab
import numpy as np

# 加载 MNIST 数据集
mnist = tf.keras.datasets.mnist
# 如果文件不存在，函数会自动下载数据集。这里使用的时绝对路径
(train_images, train_labels), (test_images, test_labels) = mnist.load_data(path='E://py/aipy/3/4/mnist/mnist.npz')

# 初始化模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, (5, 5), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
])

# 编译模型
model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型
model.fit(train_images.reshape(-1, 28, 28, 1) / 255.0, train_labels, epochs=5)
# 再次训练模型
def train_model():
    global model
    model.fit(train_images.reshape(-1, 28, 28, 1) / 255.0, train_labels, epochs=5)
    print("模型训练完成！")

# 测试模型
def test_model():
    global model
    test_loss, test_acc = model.evaluate(test_images.reshape(-1, 28, 28, 1) / 255.0, test_labels)
    print("测试准确率:", test_acc)

# 识别数字
def recognize_digit():
    global model
    # 获取 Canvas 的坐标
    x =  canvas.winfo_rootx()
    y =  canvas.winfo_rooty()
    # 使用 ImageGrab 捕获 Canvas 的内容
    img = ImageGrab.grab((x, y, x + w, y + h))
    # 缩放图像至 28x28 像素
    img = img.resize((28, 28))
    # img.save("3/4/mnist/temp.png")
    # 转换为灰度图像
    img = img.convert('L')
    image_array = np.array(img)
    # 归一化
    image_array = image_array / 255.0
    # 预测
    prediction = model.predict(image_array.reshape(1, 28, 28, 1))
    # 显示结果
    result = np.argmax(prediction)
    print("识别结果:", result)

# 初始化画布
def clear_canvas():
    canvas.delete("all")

# 绘制线条
def draw_line(event):
    x, y = event.x, event.y
    canvas.create_line(x, y, x+2, y+2, fill='white', width=20)

# 创建 Tkinter 窗口
root = tk.Tk()
root.title("MNIST 手写数字识别")
# 画布大小
w = 280
h = 280
# 创建画布
canvas = tk.Canvas(root, width=w, height=h, bg='black')
canvas.pack()
# 绑定鼠标事件
canvas.bind("<B1-Motion>", draw_line)

# 创建按钮
tk.Button(root, text="训练模型", command=train_model).pack(side=tk.LEFT)
tk.Button(root, text="测试模型", command=test_model).pack(side=tk.LEFT)
tk.Button(root, text="识别数字", command=recognize_digit).pack(side=tk.LEFT)
tk.Button(root, text="清空画布", command=clear_canvas).pack(side=tk.LEFT)

# 运行 Tkinter 主循环
root.mainloop()